#include "cipher.h"
#include "caes.h"
#include "zbase64.h"
#include "RSA.h"

//get the bit length after zero padding (one of 128, 192, 256), skey.length() must be <= 32.
int keybits(const std::string &skey)
{
    int len = skey.length() < 16 ? 16 : skey.length(); //make sure (len/8)*64 >= 128
    return (len / 8) * 64;
}

//a向上取整为b的倍数
inline int roundon(int a, int b)
{
    return a % b == 0 ? a : a + b - a % b;
}

//copy and do Pkcs7 padding, src needs to be null-terminated, buff needs to be big enough.
int padcpy(uint8_t *buff, const void *src)
{
    auto srclen = strlen((const char *)src);
    memcpy(buff, src, srclen);

    uint8_t padval = BLOCK_SIZE - srclen % BLOCK_SIZE;  //0 < padval <= 16
    for (size_t i = srclen; i < srclen + padval; i++)
        buff[i] = padval;
    return srclen + padval;
}

std::pair<std::string, std::string> pwd_encrypt(std::string strText, std::string raw_key){

    std::string general_key;                 /*声明用来保存随机字符串的str*/
    char c;                     /*声明字符c，用来保存随机生成的字符*/
    int idx;                    /*用来循环的变量*/
    /*循环向字符串中添加随机生成的字符*/
    for(idx = 0;idx < 16;idx ++)
    {
        /*rand()%26是取余，余数为0~25加上'a',就是字母a~z,详见asc码表*/
        c = 'a' + rand()%26;
        general_key.push_back(c);       /*push_back()是string类尾插函数。这里插入随机字符c*/
    }

    std::string aes_key= dakuang::CBASE64::encode(general_key);
    std::string encryptedText = aes_encrypt_string(strText, general_key);

    std::string strPublicKey = "-----BEGIN RSA PUBLIC KEY-----\n" + raw_key + "\n-----END RSA PUBLIC KEY-----\n";
    std::string strCipher = "";
    std::string keyCipher = "";

    dakuang::CRSA::encryptByPublicPEMString(aes_key, keyCipher, strPublicKey);
    dakuang::CRSA::encryptByPublicPEMString(encryptedText, strCipher, strPublicKey);

    return {dakuang::CBASE64::encode(keyCipher), dakuang::CBASE64::encode(strCipher)};
}

//note: the length of skey should be <= 32; no length check conducted
std::string aes_encrypt_string(const std::string &plaintext, const std::string &skey)
{
    aes_context ctx;
    uint8 key[32] = {0}; //zero padding
    memcpy(key, skey.c_str(), skey.length());
    aes_set_key(&ctx, key, keybits(skey));
    size_t outlen = (plaintext.length() / BLOCK_SIZE + 1) * BLOCK_SIZE;	//padding之后plaintext的长度
    uint8 *output = new uint8[outlen];
    padcpy(output, plaintext.c_str());

    //每次处理16字节
    for (size_t i = 0; i < outlen; i += BLOCK_SIZE)
        aes_encrypt(&ctx, output+i, output+i); //note that input and output could use the same buffer

    size_t baselen = roundon(outlen, 3) / 3 * 4;
    char *basebuf = new char[baselen];
    base64_encode(output, outlen, basebuf);
    std::string ret(basebuf, baselen);
    delete[] basebuf;
    delete[] output;
    return ret;
}

std::string aes_decrypt_string(const std::string &basestr, const std::string &skey)
{
    aes_context ctx;
    uint8 key[32] = {0};
    memcpy(key, skey.c_str(), skey.length());
    aes_set_key(&ctx, key, keybits(skey));

    size_t outlen = basestr.length() / 4 * 3;
    uint8 *output = new uint8[outlen]; //because of the trailing '='s, there may be one or two extra bytes, no biggie
    outlen = base64_decode(basestr.c_str(), basestr.length(), output);

    //每次处理16字节
    for (size_t i = 0; i < outlen; i += BLOCK_SIZE)
        aes_decrypt(&ctx, output+i, output+i);

    output[outlen - output[outlen-1]] = '\0';	//unpad
    std::string ret = (char *)output;
    delete[] output;
    return ret;
}
